import cv2
import json
import numpy as np
import torch
from skimage import io,transform,color
from torchvision import transforms, models
# 自己的模型

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

num_classes = 146
MAX_EPOCH = 6
BATCH_SIZE = 8
LR = 0.002
log_interval = 10
start_epoch = -1
lr_decay_step = 7

train_dir = '../../../../home/data/'
num = os.listdir(train_dir)
train_dir = os.path.join(train_dir,num[0])
# 仅用于编码测试
type_lst = os.listdir(train_dir)
classes = type_lst
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]

val_transform = transforms.Compose([
    transforms.Resize((112,112)),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])
# 构建MyDataset实例
val_data = TrashDataset(data_dir = train_dir,transform=train_transform)
# 构建DataLoder
val_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)


def init():
    # 测试时选择的文件名
    pth = '../models/202158model_resnet_20epoch/models.pkl'
    model = torch.load(pth) 
    
    return model

# 根据训练的标签设置
class_dict = {}
f = open('./class.txt','r')
a = f.read()
class_dict = eval(a)
class_dict = {value:key for key, value in class_dict.items()}
f.close()

def process_image(net, input_image, args=None):
    
    img = cv2.resize(input_image,(112,112)).astype(np.float32)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    # h, w, c to c, h, w
    img = img.transpose(2,0,1)
    img = cv2.normalize(img,None)
    img = torch.tensor([img]).cuda()
#     torch.set_grad_enabled(False)
    net.eval()
    with torch.no_grad():
        out = net(img)
        print(out)
        _, pred = torch.max(out.data, 1)
        data = json.dumps({'class': class_dict[pred[0].item()]},indent=4)
    return data, pred


if __name__ == '__main__':
    net = init()
    for j, data in enumerate(valid_loader):
        inputs, labels = data
        print(process_image(net, inputs))
        dic, predicted = process_image(net, inputs)
        loss = criterion(outputs, labels)

        total_val += labels.size(0)
        correct_val += (predicted == labels).squeeze().cpu().sum().numpy()

        loss_val += loss.item()

        loss_val_mean = loss_val/len(valid_loader)
        valid_curve.append(loss_val_mean)
        print("Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
            epoch, MAX_EPOCH, j+1, len(valid_loader), loss_val_mean, correct_val / total_val))

#     x = torch.randn((112, 112, 3)).numpy()

